#!/usr/bin/env python3

import argparse
import copy
import os
from contextlib import contextmanager
from pathlib import Path

from configs.util import load_super_config, update_config
from submit_and_evaluate import submit_and_evaluate
from utils.misc import yaml_to_dict


def load_config(yaml_path: str) -> dict:
    cfg = yaml_to_dict(yaml_path)
    super_path = cfg.get("SUPER_CONFIG_PATH")
    if super_path:
        parent = load_super_config(super_path)
        cfg = update_config(parent, cfg)
    return cfg


def prepare_common(cfg: dict, checkpoint: str, output_root: Path, data_root: str | None) -> dict:
    prepared = copy.deepcopy(cfg)
    prepared["INFERENCE_MODEL"] = checkpoint
    prepared["OUTPUTS_DIR"] = str(output_root)
    if data_root:
        prepared["DATA_ROOT"] = data_root
    if prepared.get("INFERENCE_MODE") is None:
        prepared["INFERENCE_MODE"] = "evaluate"
    if prepared.get("EXP_NAME") is None:
        prepared["EXP_NAME"] = output_root.name
    return prepared


@contextmanager
def _temporary_environ(overrides: dict[str, str]) -> None:
    saved: dict[str, str | None] = {}
    try:
        for key, value in overrides.items():
            saved[key] = os.environ.get(key)
            if value is None:
                os.environ.pop(key, None)
            else:
                os.environ[key] = value
        yield
    finally:
        for key, previous in saved.items():
            if previous is None:
                os.environ.pop(key, None)
            else:
                os.environ[key] = previous


def run_once(base_cfg: dict, *, tag: str, use_lwg: bool, lwg_model: str | None) -> None:
    cfg = copy.deepcopy(base_cfg)
    cfg["OUTPUTS_DIR"] = str(Path(base_cfg["OUTPUTS_DIR"]) / tag)
    cfg["INFERENCE_GROUP"] = f"{cfg.get('INFERENCE_GROUP', 'default')}_{tag}"
    cfg["USE_LWG"] = use_lwg
    use_buffer_gate = cfg.get("USE_BUFFER_GATE", False)
    if use_lwg:
        if not lwg_model:
            raise ValueError("LWG run requested but --lwg-model not provided.")
        cfg["LWG_MODEL_PATH"] = lwg_model
        env_overrides = {
            "USE_LWG": "True",
            "USE_BUFFER_GATE": "True" if use_buffer_gate else "False",
        }
    else:
        cfg["EXPORT_LWG_FEATURES"] = False
        cfg["USE_BUFFER_GATE"] = False
        env_overrides = {
            "USE_LWG": "False",
            "USE_BUFFER_GATE": "False",
        }
    with _temporary_environ(env_overrides):
        submit_and_evaluate(cfg)


def main() -> None:
    parser = argparse.ArgumentParser("Compare MOTIP baseline vs MA-MOTIP LWG inference.")
    parser.add_argument("--config", required=True, help="YAML config path for MA-MOTIP.")
    parser.add_argument("--checkpoint", required=True, help="Baseline checkpoint .pth path.")
    parser.add_argument("--output-root", default="./outputs/lwg_compare", help="Root directory for comparison outputs.")
    parser.add_argument("--data-root", help="Override DATA_ROOT for evaluation runs.")
    parser.add_argument("--lwg-model", help="Trained LWG model path for the LWG-enabled run.")
    parser.add_argument("--skip-lwg", action="store_true", help="Only run baseline without LWG.")
    args = parser.parse_args()

    cfg = load_config(args.config)
    output_root = Path(args.output_root).resolve()
    base_cfg = prepare_common(cfg, args.checkpoint, output_root, args.data_root)

    run_once(base_cfg, tag="baseline", use_lwg=False, lwg_model=None)

    if not args.skip_lwg:
        run_once(base_cfg, tag="lwg", use_lwg=True, lwg_model=args.lwg_model)


if __name__ == "__main__":
    main()
